#! /usr/bin/python
# encoding=utf-8
# Created by Fenglu Niu on 2025/5/6 14:17
import json

from langchain_core.messages import AIMessage, ToolMessage, AIMessageChunk
from langchain_mcp_adapters.client import MultiServerMCPClient
from langgraph.prebuilt import create_react_agent

import common
from api.entity import Message
from api.enums import EventType
from api.service.message_service import MessageService
from common import REDIS_KEY_CONVERSATION_AGENTS
from common.mysql_client import MysqlClient
from common.redis_client import RedisClient
from common.tencent_tts_api import TencentTTSApi
from hk_mcp import mcp_tools_desc


async def mcp_client(conversation_id, msgs, model):
    yield {
        "event": EventType.START.value,
        "data": json.dumps({'conversation_id': conversation_id})
    }
    async with MultiServerMCPClient({
        'school_life': {
            'url': 'http://localhost:8000/sse',
            'transport': 'sse',
        },
        # 'amap': {
        #
        # },
        # 'wanx': {
        #
        # },
        # 'tongxiao': {
        #
        # }
    }) as client:
        tools = client.get_tools()
        # print(tools)
        # Create and run the agent
        agent = create_react_agent(model, tools=tools)
        response = agent.astream({"messages": msgs}, stream_mode="messages")
        hk_messages = []
        prev_event = ''
        prev_data = {}
        llm_generation = ''
        prev_voice_text = ''
        prev_voice_char_count = 0
        voice_max_char_count = 77
        agents = json.loads(await RedisClient().get_aredis().get(REDIS_KEY_CONVERSATION_AGENTS.format(conversation_id)))
        async for messages in response:
            # print(messages)
            message = messages[0]
            if isinstance(message, AIMessageChunk):
                if len(message.response_metadata) > 0 and message.response_metadata['finish_reason'] == 'tool_calls':
                    if prev_event != EventType.FUNCTION.value and prev_event != EventType.FUNCTION_START.value:
                        common.get_logger().error('出错:当出现此消息时，上一条消息必须是FUNCTION或者FUNCTION_START')
                        yield {
                            "event": EventType.ERROR.value,
                            "data": "出错:当出现此消息时，上一条消息必须是FUNCTION或者FUNCTION_START"
                        }
                        break
                    else:
                        data = json.dumps(prev_data, ensure_ascii=False)
                        yield {
                            "event": prev_event,
                            "data": data
                        }
                        yield {
                            "event": EventType.FUNCTION_END.value,
                            "data": EventType.FUNCTION_END.value
                        }
                        hk_messages.append(Message(conversation_id=conversation_id, role='ai', content=data, tag='tool_calls'))
                        prev_event = ''
                        prev_data = {}

                if len(message.response_metadata) > 0 and message.response_metadata['finish_reason'] == 'stop':
                    data = message.content.replace('<think>\n\n</think>\n\n', '')
                    if data and not data.isspace():
                        prev_voice_text += data
                        prev_voice_char_count += len(data)
                        llm_generation += data
                        data = json.dumps(data, ensure_ascii=False)
                        yield {
                            "event": EventType.MESSAGE.value,
                            "data": data
                        }
                        prev_event = ''
                        prev_data = {}
                    hk_messages.append(
                        Message(conversation_id=conversation_id, role='ai', content=llm_generation, tag='history'))
                    if prev_voice_char_count > 0:
                        try:
                            data = TencentTTSApi().tts(prev_voice_text, agents)
                            yield {
                                "event": EventType.VOICE.value,
                                "data": json.dumps(data, ensure_ascii=False)
                            }
                        except Exception as e:
                            print(e)
                            yield {
                                'event': EventType.VOICE_ERROR.value,
                                'data': EventType.VOICE_ERROR.value
                            }
                        prev_voice_text = ''
                        prev_voice_char_count = 0
                    continue

                if not message.content and len(message.tool_call_chunks) == 0:
                    continue
                if not message.content and len(message.tool_call_chunks) > 0:
                    function_name = message.tool_call_chunks[0]['name']
                    args = message.tool_call_chunks[0]['args']
                    if not prev_event:
                        prev_event = EventType.FUNCTION_START.value
                        yield {
                            "event": prev_event,
                            "data": EventType.FUNCTION_START.value
                        }
                        prev_event = EventType.FUNCTION.value
                        function_description = mcp_tools_desc[function_name]
                        prev_data['name'] = function_name
                        prev_data['description'] = function_description
                        prev_data['id'] = message.tool_call_chunks[0]['id']
                        prev_data['args'] = (args if args else '')
                    else:
                        if function_name:
                            prev_event = EventType.FUNCTION.value
                            if len(prev_data) > 0:
                                data = json.dumps(prev_data, ensure_ascii=False)
                                yield {
                                    "event": prev_event,
                                    "data": data
                                }
                                hk_messages.append(Message(conversation_id=conversation_id, role='ai', content=data, tag='tool_calls'))
                                prev_data = {}
                            function_description = mcp_tools_desc[function_name]
                            prev_data['name'] = function_name
                            prev_data['description'] = function_description
                            prev_data['id'] = message.tool_call_chunks[0]['id']
                            prev_data['args'] = args if args else ''
                        else:
                            prev_data['args'] = prev_data['args'] + (args if args else '')
                else:
                    if not prev_event and message.content.isspace():
                        continue
                    if prev_event != EventType.MESSAGE.value:
                        prev_data = {}
                        llm_generation = ''

                    prev_event = EventType.MESSAGE.value
                    data = message.content.replace('<think>\n\n</think>\n\n', '')
                    prev_voice_text += data
                    prev_voice_char_count += len(data)
                    llm_generation += data
                    data = json.dumps(data, ensure_ascii=False)
                    yield {
                        "event": EventType.MESSAGE.value,
                        "data": data
                    }
                    if prev_voice_char_count >= voice_max_char_count:
                        try:
                            data = TencentTTSApi().tts(prev_voice_text, agents)
                            yield {
                                "event": EventType.VOICE.value,
                                "data": json.dumps(data, ensure_ascii=False)
                            }
                        except Exception as e:
                            print(e)
                            yield {
                                'event': EventType.VOICE_ERROR.value,
                                'data': EventType.VOICE_ERROR.value
                            }
                        prev_voice_text = ''
                        prev_voice_char_count = 0

            elif isinstance(message, ToolMessage):
                prev_event = EventType.FUNCTION_RESP.value
                prev_data = {'content': message.content, 'tool_call_id': message.tool_call_id, 'name': message.name}
                data = json.dumps(prev_data, ensure_ascii=False)
                hk_messages.append(Message(conversation_id=conversation_id, role='tool', content=data))
                yield {
                    "event": prev_event,
                    "data": data
                }
                prev_event = ''
                prev_data = {}

        MessageService.create_messages(hk_messages, next(MysqlClient().get_session()))
        yield {
            "event": EventType.END.value,
            "data": "end"
        }
